from typing import List

import numpy as np
from langchain.embeddings.base import Embeddings

from vec_database.config import EMBEDDING_MODEL
from vec_database.embeddings_tool.embeddings_api import embed_texts, aembed_texts


class EmbeddingsFunAdapter(Embeddings):
    def __init__(self, embed_model: str = EMBEDDING_MODEL):
        self.embed_model = embed_model

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        embeddings = embed_texts(texts=texts, embed_model=self.embed_model, to_query=False).data
        return normalize(embeddings).tolist()

    def embed_query(self, text: str) -> List[float]:
        embeddings = embed_texts(texts=[text], embed_model=self.embed_model, to_query=True).data
        query_embed = embeddings[0]
        query_embed_2d = np.reshape(query_embed, (1, -1))  # 将一维数组转换为二维数组
        normalized_query_embed = normalize(query_embed_2d)
        return normalized_query_embed[0].tolist()  # 将结果转换为一维数组并返回

    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
        embeddings = (await aembed_texts(texts=texts, embed_model=self.embed_model, to_query=False)).data
        return normalize(embeddings).tolist()

    async def aembed_query(self, text: str) -> List[float]:
        embeddings = (await aembed_texts(texts=[text], embed_model=self.embed_model, to_query=True)).data
        query_embed = embeddings[0]
        query_embed_2d = np.reshape(query_embed, (1, -1))  # 将一维数组转换为二维数组
        normalized_query_embed = normalize(query_embed_2d)
        return normalized_query_embed[0].tolist()  # 将结果转换为一维数组并返回


# sklearn.preprocessing.normalize 的替代（使用 L2），避免安装 scipy, scikit-learn
def normalize(embeddings: List[List[float]]) -> np.ndarray:
    # np.linalg.norm() 是 NumPy 库中的一个函数，用于计算向量或矩阵的范数
    norm = np.linalg.norm(embeddings, axis=1)
    # np.reshape(a,(m*n))，按给定索引顺序展开数组a，将a重构成m*n的数组
    norm = np.reshape(norm, (norm.shape[0], 1))
    norm = np.tile(norm, (1, len(embeddings[0])))
    # numpy.divide()函数基于NumPy库的高性能数组操作机制，实现了向量和矩阵的除法运算
    return np.divide(embeddings, norm)

'''
np.linalg.norm() 是 NumPy 库中的一个函数，用于计算向量或矩阵的范数。范数是一种衡量向量或矩阵大小的度量方式，在线性代数中有广泛的应用。这个函数可以接受多个参数，主要有以下两个参数：1

arr：表示输入的向量或矩阵。可以是一维数组（向量）或多维数组（矩阵）。
ord：表示计算范数的类型。常用的取值有：
ord=None（默认值）：计算 Frobenius 范数，对于矩阵，即各个元素的平方和的平方根；对于向量，即向量元素的平方和的平方根。
ord=1：计算 L1 范数，对于矩阵，即各列元素绝对值之和的最大值；对于向量，即向量元素绝对值之和。
ord=2：计算 L2 范数（默认），对于矩阵，即特征值的平方和的平方根（即谱范数）；对于向量，即向量元素的平方和的平方根。
'''

'''
如果参数只有一个是按照X轴计算
np.tile(a,(2,1)), 沿Y轴扩大2倍，沿X轴扩大1倍
np.title(a,(2)),沿X轴扩大2倍
'''